#!/opt/local/bin/python

from rfvectors import cavity, scale_roverq
from diffeqns import transientcavity, analysis, calcCurrents
import csv
from scipy import nan, exp, sqrt, arctan, pi, isnan
from numpy import array
import matplotlib.pyplot as plt
rad2deg = lambda x: x*180/pi

class linacsection(object):
    def __init__(self, fname, RoverQ, QL, freq, Ib0, beta0, Eampfile=nan):
        self.fname = fname
        self.RoverQ = RoverQ
        self.QL = QL
        self.freq = freq
        self.Ib0 = Ib0
        self.beta0 = beta0
        self.Eampfile = Eampfile
        self.savePlots = False
        self.plotstart = 0

        self.getData()

    def plots(self, start=0):
        if self.savePlots:
            self.savePlots = False
            self.plotstart = start
        else:
            self.savePlots = True
            self.plotstart = start

    def getData(self):
        self.data = csv.reader(open(self.fname, 'rU'), dialect=csv.excel_tab, delimiter=',')

    def integRefl(self, I_b, pulselen, acc, faulttime=nan):
        if isnan(faulttime):
            faulttime = pulselen
        integPref = []
        count = 0
        for row in self.data:
            count += 1
            print "Cavity number = %d" % count
            V = float(row[0]) * float(row[3])
            beamph = float(row[1])
            beta = float(row[2])

            sscav = cavity(f=self.freq, delF=nan, I_b0=self.Ib0, I_b=I_b, V_cav=V, beamph=beamph,
                    RoverQ=self.RoverQ, Ql=self.QL, beta_0=self.beta0, beta=beta, Eampfile=self.Eampfile)
            RoverQscaling = scale_roverq(beta_0=sscav.beta_0, file=sscav.Eampfile)(beta=sscav.beta, freq=sscav.f)
            try:
                RoverQscaling = RoverQscaling[0]
            except IndexError:
                pass
            tcav = transientcavity(freq=sscav.f, QL=sscav.Ql, RoverQ=sscav.RoverQ, detuning=sscav.delF)
            cavanal = analysis(tcavobj=tcav, timestop=5e-3, acc=acc)

            Ig, Ibeam, tinj = calcCurrents(cavobj=sscav, tcavobj=tcav,
                    time=cavanal.time, pulselength=pulselen, RoverQscaling=RoverQscaling)
            for i in range(len(cavanal.time)):
                if cavanal.time[i] < tinj + faulttime: continue
                Ibeam[i] = 0 + 0j
            I= Ig + Ibeam

            print "Running......"
            cavanal.run(iniV=0, I=I)
            print "Done!"
            Vinc  = 0.5 * Ig * exp(tcav.phi*1j) * tcav.RL
            Pinc  = abs(Vinc)**2 / tcav.RL
            Vref  = abs(Vinc) - sqrt(cavanal.V[:,0]**2 + cavanal.V[:,1]**2)
            Pref  = Vref**2 / tcav.RL
            Vbeam = 0.5 * Ibeam * exp(tcav.phi*1j) * tcav.RL
            Pbeam = abs(Vbeam)**2 / tcav.RL

            if self.savePlots:
                fig2 = plt.figure((self.plotstart+count)*1e3 + 2)
                ax1fig2 = fig2.add_subplot(211)
                ax1fig2.plot(1e3*cavanal.time, 1e-6*Vinc.real, 'b', label='Real voltage')
                ax1fig2.plot(1e3*cavanal.time, 1e-6*Vinc.imag, 'r', label='Imag voltage')
                ax1fig2.plot(1e3*cavanal.time, 1e-6*abs(Vbeam), 'k', label='Abs beam voltage')
                ax1fig2.plot(1e3*cavanal.time, 1e-6*abs(Vinc), 'g', label='Abs voltage')
                ax1fig2.set_ylabel('Voltage / MV')
                
                ax2fig2 = fig2.add_subplot(212)
                ax2fig2.plot(1e3*cavanal.time, abs(Vref)/abs(Vinc), label='Reflection coeff.')
                ax2fig2.set_ylabel('Reflection coefficient')
                
                fig1 = plt.figure((self.plotstart+count)*1e3 + 1)
                ax1 = fig1.add_subplot(211)
                ax2 = fig1.add_subplot(212)
                ax1.plot(1e3*cavanal.time, 1e-6*cavanal.V[:,0], '-', label='Real cavity voltage')
                ax1.plot(1e3*cavanal.time, 1e-6*cavanal.V[:,1], '-', label='Imag cavity voltage')
                ax2.plot(1e3*cavanal.time, 1e-6*sqrt(cavanal.V[:,0]**2 + cavanal.V[:,1]**2), '-', label='Absolute cavity voltage')
                ax2.plot(1e3*cavanal.time, 1e-6*Vref, 'r', label='Reflected voltage')
                
                fig3 = plt.figure((self.plotstart+count)*1e3 + 3)
                ax1fig3 = fig3.add_subplot(111)
                ax1fig3.plot(1e3*cavanal.time, rad2deg(arctan(cavanal.V[:,1]/cavanal.V[:,0])), label='RF phase')
                ax1fig3.set_ylabel('RF phase (degrees)')
                
                fig4 = plt.figure((self.plotstart+count)*1e3 + 4)
                ax1fig4 = fig4.add_subplot(111)
                ax1fig4.plot(1e3*cavanal.time, Pref/1e3, label='Reflected: %0.3f J' % (sum(Pref * cavanal.timestep)))
                ax1fig4.set_ylabel('Power (kW)')
                
                for a in [ax1,ax2,ax1fig2,ax2fig2,ax1fig3,ax1fig4]:
                    a.grid()
                    a.set_xlabel('Time / ms')
                    a.set_xlim(right=1e3*cavanal.time[-1])
                    a.axvline(x=tinj*1e3, color='g')
                    a.axvline(x=(tinj+pulselen)*1e3, color='g')
                    a.legend(loc=0)
                ax1fig2.legend(loc=1)
                ax1.set_ylabel('Voltage / MV')
                ax2.set_ylabel('Abs voltage / MV')

                fig1.savefig('individualcavs/cav%d_cavvoltages.pdf' % (self.plotstart+count))
                fig2.savefig('individualcavs/cav%d_incidentvoltages.pdf' % (self.plotstart+count))
                fig3.savefig('individualcavs/cav%d_cavphase.pdf' % (self.plotstart+count))
                fig4.savefig('individualcavs/cav%d_reflectedpower.pdf' % (self.plotstart+count))

            integPref.append( sum(Pref * cavanal.timestep) )

        return integPref

if __name__ == "__main__":
    opts = {"I_b": 1e-3, "pulselen": 2.8e-3, "faulttime":1e-6, "acc": 1e-4}
    accumcav = 0
    spokePref = []
    Eref = []
    A = linacsection(fname='spokenew_sec1.csv', RoverQ=500, QL=4.287e5, freq = 352.21e6, Ib0=50e-3, beta0=0.46, Eampfile='Ezvszspoke.dat')
    #A.plots(start=accumcav)
    integPref = A.integRefl(**opts)
    for i in integPref:
        spokePref.append( i )
    for i in integPref:
        Eref.append( i )
    accumcav += len(integPref)

    A = linacsection(fname='spokenew_sec2.csv', RoverQ=500, QL=2.403e5, freq = 352.21e6, Ib0=50e-3, beta0=0.46, Eampfile='Ezvszspoke.dat')
    #A.plots(start=accumcav)
    integPref = A.integRefl(**opts)
    for i in integPref:
        spokePref.append( i )
    for i in integPref:
        Eref.append( i )
    accumcav += len(integPref)

    A = linacsection(fname='spokenew_sec3.csv', RoverQ=500, QL=1.500e5, freq = 352.21e6, Ib0=50e-3, beta0=0.46, Eampfile='Ezvszspoke.dat')
    #A.plots(start=accumcav)
    integPref = A.integRefl(**opts)
    for i in integPref:
        spokePref.append( i )
    for i in integPref:
        Eref.append( i )
    accumcav += len(integPref)

    #fig1 = plt.figure(1)
    #ax1fig1 = fig1.add_subplot(111)
    #ax1fig1.plot(spokePref, 'x')
    #ax1fig1.grid()
    #ax1fig1.set_xlabel('Cavity number')
    #ax1fig1.set_ylabel('Reflected energy (J)')
    #ax1fig1.set_ylim(bottom=0)

    A = linacsection(fname='lowbetanew.csv', RoverQ=300, QL=9.176e5, freq = 704.42e6, Ib0=50e-3, beta0=0.7, Eampfile=nan)
    #A.plots(start=accumcav)
    integPref = A.integRefl(**opts)
    for i in integPref:
        Eref.append( i )
    accumcav += len(integPref)

    #fig2 = plt.figure(2)
    #ax1fig2 = fig2.add_subplot(111)
    #ax1fig2.plot(array(integPref), 'x')
    #ax1fig2.grid()
    #ax1fig2.set_xlabel('Cavity number')
    #ax1fig2.set_ylabel('Reflected energy (J)')
    #ax1fig2.set_ylim(bottom=0)

    A = linacsection(fname='highbetanew.csv', RoverQ=477, QL=8.123e5, freq = 704.42e6, Ib0=50e-3, beta0=0.92, Eampfile=nan)
    #A.plots(start=accumcav)
    integPref = A.integRefl(**opts)
    for i in integPref:
        Eref.append( i )
    accumcav += len(integPref)

    #fig3 = plt.figure(3)
    #ax1fig3 = fig3.add_subplot(111)
    #ax1fig3.plot(array(integPref), 'x')
    #ax1fig3.grid()
    #ax1fig3.set_xlabel('Cavity number')
    #ax1fig3.set_ylabel('Reflected energy (J)')
    #ax1fig3.set_ylim(bottom=0)

    fig10 = plt.figure(10)
    ax1fig10 = fig10.add_subplot(111)
    ax1fig10.plot(Eref, 'x')
    ax1fig10.grid()
    ax1fig10.set_xlabel('Cavity number')
    ax1fig10.set_ylabel('Reflected energy (J)')
    ax1fig10.set_ylim(bottom=0)

    plt.show()

